--WAF Action
require 'config'
require 'lib'

--args
local rulematch = ngx.re.find
local unescape = ngx.unescape_uri

--allow white ip
function white_ip_check()
     if config_white_ip_check == "on" then
        local IP_WHITE_RULE = get_rule('whiteip.rule')
        local WHITE_IP = get_client_ip()
        if IP_WHITE_RULE ~= nil then
            for _,rule in pairs(IP_WHITE_RULE) do
                if rule ~= "" and rulematch(WHITE_IP,rule,"jo") then
                    return true
                end
            end
        end
    end
end

--deny black ip
function black_ip_check()
    if config_black_ip_check == "on" then
       local IP_BLACK_RULE = get_rule('blackip.rule')
       local BLACK_IP = get_client_ip()
       if IP_BLACK_RULE ~= nil then
           for _,rule in pairs(IP_BLACK_RULE) do
               if rule ~= "" and rulematch(BLACK_IP,rule,"jo") then
                    if config_waf_enable == "on" then
                        log_record('Deny_IP',ngx.var_request_uri,"_","_")
                        waf_output()
                        return true
                    end
               end
           end
       end
   end
end

--deny user agent
function user_agent_attack_check()
    if config_user_agent_check == "on" then
        local USER_AGENT_RULES = get_rule('useragent.rule')
        local USER_AGENT = ngx.var.http_user_agent
        if USER_AGENT_RULES ~= nil and USER_AGENT ~= nil then
            for _,rule in pairs(USER_AGENT_RULES) do
                if rule ~="" and rulematch(USER_AGENT,rule,"jo") then
                    if config_waf_enable == "on" then
                        log_record('Deny_USER_AGENT',ngx.var.request_uri,"-",rule)
                        waf_output()
                        return true
                    end
                end
            end
        end
    end
end

--deny cc attack
function cc_attack_check()
    if config_cc_check == "on" then
        local USER_AGENT = get_user_agent()
        --local ATTACK_URL = ngx.var.host .. ngx.var.request_uri
        local ATTACK_URL = ngx.var.host .. ngx.var.uri
        local CC_TOKEN = get_client_ip() .. "." .. ngx.md5(string.lower(ATTACK_URL) .. USER_AGENT)
        local limit = ngx.shared.limit
        local CCcount=tonumber(string.match(config_cc_rate,'(.*)/'))
        local CCseconds=tonumber(string.match(config_cc_rate,'/(.*)'))
        local req,_ = limit:get(CC_TOKEN)
        if req then
            -- write('/data/wwwlogs/info.log',CC_TOKEN ..'\t'.. ATTACK_URL .. '\t'.. 'req: ' .. req .. "\n")
            if req > CCcount then
                if config_waf_enable == "on" then
                    log_record('CC_Attack',ngx.var.request_uri,"-","-")
                    waf_output()
                end
            else
                limit:incr(CC_TOKEN,1)
            end
        else
            limit:set(CC_TOKEN,1,CCseconds)
        end
    end
end

--deny cookie
function cookie_attack_check()
    if config_cookie_check == "on" then
        local COOKIE_RULES = get_rule('cookie.rule')
        local USER_COOKIE = ngx.var.http_cookie
        if COOKIE_RULES ~= nil and USER_COOKIE ~= nil then
            for _,rule in pairs(COOKIE_RULES) do
                if rule ~="" and rulematch(USER_COOKIE,rule,"jo") then
                    if config_waf_enable == "on" then
                        log_record('Deny_Cookie',ngx.var.request_uri,"-",rule)
                        waf_output()
                        return true
                    end
                end
             end
	    end
    end
end

--allow white url
function white_url_check()
    if config_white_url_check == "on" then
        local URL_WHITE_RULES = get_rule('whiteurl.rule')
        local REQ_URI = string.lower(ngx.var.uri)
        if URL_WHITE_RULES ~= nil then
            for _,rule in pairs(URL_WHITE_RULES) do
                if rule ~= "" and rulematch(REQ_URI,rule,"jo") then
                    return true
                end
            end
        end
    end
end

--deny black url
function url_attack_check()
    if config_url_check == "on" then
        local URL_RULES = get_rule('blackurl.rule')
        local REQ_URI = string.lower(ngx.var.uri)
        if URL_RULES ~= nil then
            for _,rule in pairs(URL_RULES) do
                if rule ~="" and rulematch(REQ_URI,rule,"jo") then
                    if config_waf_enable == "on" then
                        log_record('Deny_URL',REQ_URI,"-",rule)
                        waf_output()
                        return true
                    end
                end
            end
        end
    end
end

--deny url args
function url_args_attack_check()
    if config_url_args_check == "on" then
        local ARGS_RULES = get_rule('args.rule')
        if URL_RULES ~= nil then
            for _,rule in pairs(ARGS_RULES) do
                local REQ_ARGS = ngx.req.get_uri_args()
                for key, val in pairs(REQ_ARGS) do
                    if type(val) == 'table' then
                        local ARGS_DATA = table.concat(val, " ")
                    else
                        local ARGS_DATA = val
                    end
                    if ARGS_DATA and type(ARGS_DATA) ~= "boolean" and rule ~="" and rulematch(unescape(ARGS_DATA),rule,"jo") then
                        if config_waf_enable == "on" then
                            log_record('Deny_URL_Args', ngx.var.request_uri, ARGS_DATA ,rule)
                            waf_output()
                            return true
                        end
                    end
                end
            end
        end
    end
end

--deny post
function post_attack_check()
    if config_post_check == "on" then
        local POST_RULES = get_rule('post.rule')
        for _,rule in pairs(POST_RULES) do
            local POST_ARGS = ngx.req.get_post_args()
        end
    end
end

--限流
function current_limit_check()
    if config_waf_enable == "on" and config_current_limit_check ~= "off" then
        local LIMIT_RULES = get_rule('limit.rule')
        local REQ_URI = string.lower(ngx.var.request_uri)
        if LIMIT_RULES ~= nil then
            for _,rule in pairs(LIMIT_RULES) do
                if rule ~= "" and rulematch(REQ_URI,string.lower(rule),"jo") then
                    local COOKIE_VALUE = ngx.var['cookie_shiro.session']
                    if COOKIE_VALUE then
                        local limit = ngx.shared.limit
                        local seconds = tonumber(config_current_limit_check)
                        local TOKEN_VALUE = 'limit-' .. ngx.md5(get_client_ip() .. COOKIE_VALUE .. REQ_URI)
                        local count,_ = limit:get(TOKEN_VALUE)
                        if count then
                            log_record('Deny_CURRENT_LIMIT', ngx.var.request_uri, '-', rule)
                            ngx.exit(503)
                        else
                            limit:set(TOKEN_VALUE, 1, seconds)
                        end
                    else
                        log_record('Deny_CURRENT_LIMIT', ngx.var.request_uri, '-', 'Deny_NO_COOKIE')
                        ngx.exit(503)
                    end
                    return true
                end
            end
        end
    end
end
